// Copyright 2020 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Modification:
// 1. Change the return parameter of newNetIPSocket to map.
// 2. Remove unused fields parsed.

package internal

import (
	"bufio"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"strings"
)

const (
	// readLimit is used by io.LimitReader while reading the content of the
	// /proc/net/udp{,6} files. The number of lines inside such a file is dynamic
	// as each line represents a single used socket.
	// In theory, the number of available sockets is 65535 (2^16 - 1) per IP.
	// With e.g. 150 Byte per line and the maximum number of 65535,
	// the reader needs to handle 150 Byte * 65535 =~ 10 MB for a single IP.
	readLimit = 4294967296 // Byte -> 4 GiB
)

// This contains generic data structures for both udp and tcp sockets.
type (
	SocketKey struct {
		LocalAddr string
		LocalPort uint64
		RemAddr   string
		RemPort   uint64
	}

	// NetSocketStateMap represents the map between socket and its state.
	NetSocketStateMap map[SocketKey]string

	// netIPSocketLine represents the fields parsed from a single line
	// in /proc/net/{t,u}dp{,6}. Fields which are not used by IPSocket are skipped.
	// For the proc file format details, see https://linux.die.net/man/5/proc.
	netIPSocketLine struct {
		LocalAddr net.IP
		LocalPort uint64
		RemAddr   net.IP
		RemPort   uint64
		St        string
	}
)

func newNetIPSocket(file string) (NetSocketStateMap, error) {
	f, err := os.Open(file)
	if err != nil {
		return nil, err
	}
	defer func() {
		_ = f.Close()
	}()

	netIPSocket := make(NetSocketStateMap)

	lr := io.LimitReader(f, readLimit)
	s := bufio.NewScanner(lr)
	s.Scan() // skip first line with headers
	for s.Scan() {
		fields := strings.Fields(s.Text())
		line, err := parseNetIPSocketLine(fields)
		if err != nil {
			return nil, err
		}
		socketKey := SocketKey{
			LocalAddr: line.LocalAddr.String(),
			LocalPort: line.LocalPort,
			RemAddr:   line.RemAddr.String(),
			RemPort:   line.RemPort,
		}
		netIPSocket[socketKey] = line.St
	}
	if err := s.Err(); err != nil {
		return nil, err
	}
	return netIPSocket, nil
}

// the /proc/net/{t,u}dp{,6} files are network byte order for ipv4 and for ipv6 the address
// is four words consisting of four bytes each. In each of those four words the four bytes
// are written in reverse order.
func parseIP(hexIP string) (net.IP, error) {
	var byteIP []byte
	byteIP, err := hex.DecodeString(hexIP)
	if err != nil {
		return nil, fmt.Errorf("cannot parse address field in socket line %q", hexIP)
	}
	switch len(byteIP) {
	case 4:
		return net.IP{byteIP[3], byteIP[2], byteIP[1], byteIP[0]}, nil
	case 16:
		i := net.IP{
			byteIP[3], byteIP[2], byteIP[1], byteIP[0],
			byteIP[7], byteIP[6], byteIP[5], byteIP[4],
			byteIP[11], byteIP[10], byteIP[9], byteIP[8],
			byteIP[15], byteIP[14], byteIP[13], byteIP[12],
		}
		return i, nil
	default:
		return nil, fmt.Errorf("unable to parse IP %s", hexIP)
	}
}

// parseNetIPSocketLine parses a single line, represented by a list of fields.
func parseNetIPSocketLine(fields []string) (*netIPSocketLine, error) {
	line := &netIPSocketLine{}
	if len(fields) < 10 {
		return nil, fmt.Errorf(
			"cannot parse net socket line as it has less then 10 columns %q",
			strings.Join(fields, " "),
		)
	}
	var err error // parse error

	// local_address
	l := strings.Split(fields[1], ":")
	if len(l) != 2 {
		return nil, fmt.Errorf("cannot parse local_address field in socket line %q", fields[1])
	}
	if line.LocalAddr, err = parseIP(l[0]); err != nil {
		return nil, err
	}
	if line.LocalPort, err = strconv.ParseUint(l[1], 16, 64); err != nil {
		return nil, fmt.Errorf("cannot parse local_address port value in socket line: %w", err)
	}

	// remote_address
	r := strings.Split(fields[2], ":")
	if len(r) != 2 {
		return nil, fmt.Errorf("cannot parse rem_address field in socket line %q", fields[1])
	}
	if line.RemAddr, err = parseIP(r[0]); err != nil {
		return nil, err
	}
	if line.RemPort, err = strconv.ParseUint(r[1], 16, 64); err != nil {
		return nil, fmt.Errorf("cannot parse rem_address port value in socket line: %w", err)
	}

	// st
	line.St = fields[3]

	return line, nil
}
